import numpy as np

import auto_file
import gen_dataset
import test.test_datasetx as test_datasetx

# DataSetX识别一键启动脚本
if __name__ == '__main__':
    # datasetx的配置
    dtx_config = {
        'path': './'
        # 标签字典
        , 'label_dict': {
            'yes': 1,
            'no': 0
        }
        # 训练集占的比例
        , 'train_percent': 0.9
        # 图片缩放的尺寸
        , 'pic_size': 224
        , 'isLog': True
        # 训练出来的模型结果存放路径
        , 'parameter_path': './test/parameters.npy'
    }
    # 生成h5格式的数据集模型
    gen_dataset.driver(dtx_config)

    # 用生成的数据集进行学习
    path = './output/'

    train_set_x_orig, train_set_y, test_set_x_orig, test_set_y = test_datasetx.load_dataset(path)
    print(train_set_x_orig.shape)
    train_x_flatten = train_set_x_orig.reshape(train_set_x_orig.shape[0], -1).T
    test_x_flatten = test_set_x_orig.reshape(test_set_x_orig.shape[0], -1).T

    train_x = train_x_flatten / 255
    train_y = train_set_y
    test_x = test_x_flatten / 255
    test_y = test_set_y

    parameters = test_datasetx.model(train_x, train_y, keep_prob=0.75, learning_rate=0.0002
                                     , num_iterations=1000, is_plot=True)

    print("使用随机删除节点，训练集:")
    predictions_train = test_datasetx.reg_utils.predict(train_x, train_y, parameters)
    print("使用随机删除节点，测试集:")
    test_datasetx.reg_utils.predictions_test = test_datasetx.reg_utils.predict(test_x, test_y, parameters)
    # 将结果通过numpy保存参数
    np.save('parameters', parameters, allow_pickle=True)

    # 用训练出来的模型结果进行
    # 自动将图片分类的参数
    # 注意：label_dict里是数字是key,标签是value
    dtx_config = {
        'path': './origin/'
        # 标签字典
        , 'label_dict': {
            1: 'yes',
            0: 'no'
        }
        # 训练集占的比例
        , 'train_percent': 0.9
        # 图片缩放的尺寸
        , 'pic_size': 224
        , 'isLog': True
        # 训练出来的模型结果
        , 'parameter_path': './test/parameters.npy'
    }

    auto_file.pic_to_np(path=dtx_config.get('path')
                        , label_dict=dtx_config.get('label_dict')
                        , pic_size=dtx_config.get('pic_size')
                        , dtx_config=dtx_config)
